package com.relax.mybatis.encrypt.strategy;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import com.relax.mybatis.constant.Algorithm;
import com.relax.mybatis.encrypt.strategy.impl.AESEncryptStrategy;
import com.relax.mybatis.encrypt.strategy.impl.HybridAlgorithmEncryptStrategy;
import com.relax.mybatis.encrypt.strategy.impl.MD5EncryptStrategy;
import com.relax.mybatis.encrypt.strategy.impl.RSAEncryptStrategy;
import com.relax.mybatis.encrypt.strategy.impl.SM2EncryptStrategy;
import com.relax.mybatis.encrypt.strategy.impl.SM3EncryptStrategy;
import com.relax.mybatis.encrypt.strategy.impl.SM4EncryptStrategy;

/**
 * @author relax
 * @date 2022/07/05
 * @desc 加密策略上下文
 */
public class EncryptContext {

    private static final Map<Algorithm, EncryptStrategy> STRATEGY_MAP = new ConcurrentHashMap<>();

    public EncryptContext() {
        resetStrategyMap();
    }

    /**
     * 重置策略集合
     */
    private void resetStrategyMap() {
        STRATEGY_MAP.clear();
        EncryptStrategy MD5EncryptStrategy = new MD5EncryptStrategy();
        STRATEGY_MAP.put(Algorithm.MD5_32, MD5EncryptStrategy);
        STRATEGY_MAP.put(Algorithm.MD5_16, MD5EncryptStrategy);
        STRATEGY_MAP.put(Algorithm.AES, new AESEncryptStrategy());
        STRATEGY_MAP.put(Algorithm.SM2, new SM2EncryptStrategy());
        STRATEGY_MAP.put(Algorithm.SM3, new SM3EncryptStrategy());
        STRATEGY_MAP.put(Algorithm.SM4, new SM4EncryptStrategy());
        STRATEGY_MAP.put(Algorithm.RSA, new RSAEncryptStrategy());
        EncryptStrategy HybridAlgorithmEncryptStrategy = new HybridAlgorithmEncryptStrategy();
        STRATEGY_MAP.put(Algorithm.PBEWithMD5AndDES, HybridAlgorithmEncryptStrategy);
        STRATEGY_MAP.put(Algorithm.PBEWithMD5AndTripleDES, HybridAlgorithmEncryptStrategy);
        STRATEGY_MAP.put(Algorithm.PBEWithHMACSHA512AndAES_256, HybridAlgorithmEncryptStrategy);
        STRATEGY_MAP.put(Algorithm.PBEWithSHA1AndDESede, HybridAlgorithmEncryptStrategy);
        STRATEGY_MAP.put(Algorithm.PBEWithSHA1AndRC2_40, HybridAlgorithmEncryptStrategy);
    }

    /**
     * 获取策略
     * 
     * @param algorithm
     * @return
     */
    public static EncryptStrategy getStrategy(Algorithm algorithm) {
        EncryptStrategy strategy = STRATEGY_MAP.get(algorithm);
        if (strategy == null) {
            throw new IllegalArgumentException("none encrypt strategy found, algorithm: " + algorithm);
        }
        return strategy;
    }

    /**
     * 添加自定义策略
     * 
     * @param algorithm
     * @param strategy
     * @return
     */
    public static void addStrategy(Algorithm algorithm, EncryptStrategy strategy) {
        STRATEGY_MAP.put(algorithm, strategy);
    }

}
